# -*- coding: utf-8 -*-

# File Name： single_stage_detector
# Description :
# Author : lirui
# create_date： 2022/6/3
# Change Activity:

from .base_detector import BaseDetector
from ..backbone import build_backbone
from ..bbox_head import build_head
from ..neck import build_neck


class SingleStageDetector(BaseDetector):
    """
    Single stage detector base.
    """

    def __init__(self,
                 backbone_cfg,
                 neck_cfg=None,
                 bbox_head_cfg=None,
                 init_cfg=None):
        """

        Args:
            backbone_cfg: backbone_config
            neck_cfg: neck_config
            bbox_head_cfg: bbox_head_config
            init_cfg:
        """
        self.backbone = build_backbone(backbone_cfg)
        if neck_cfg is not None:
            self.neck = build_neck(neck_cfg)
        self.bbox_head = build_head(bbox_head_cfg)
        super().__init__(init_cfg)

    def extract_feat(self, x):
        """
        Extract features from images.
        """
        x = self.backbone(x)
        if self.with_neck:
            x = self.neck(x)
        return x

    def forward(self, x, img_metas):
        """
        forward function

        Args:
            x:  input
            img_metas: img info.

        Returns: result.
        """
        x = self.extract_feat(x)
        x = self.bbox_head(x, img_metas)
        return x
